package com.jakewharton.mosaic.animation

import androidx.compose.runtime.Stable
import kotlinx.coroutines.CancellationException
import kotlinx.coroutines.Job
import kotlinx.coroutines.coroutineScope
import kotlinx.coroutines.sync.Mutex
import kotlinx.coroutines.sync.withLock

// This is an internal copy of androidx.compose.foundation.MutatorMutex. Do not modify.

/**
 * Priorities for performing mutation on state.
 *
 * [MutatePriority] values follow the natural ordering of `enum class` values; a value that compares
 * as `>` to another has a higher priority. A mutation of equal or greater priority will interrupt
 * the current mutation in progress.
 */
internal enum class MutatePriority {
	/**
	 * The default priority for mutations. Can be interrupted by other [Default], [UserInput] or
	 * [PreventUserInput] priority operations. [Default] priority should be used for programmatic
	 * animations or changes that should not interrupt user input.
	 */
	Default,

	/**
	 * An elevated priority for mutations meant for implementing direct user interactions. Can be
	 * interrupted by other [UserInput] or [PreventUserInput] priority operations.
	 */
	UserInput,

	/**
	 * A high-priority mutation that can only be interrupted by other [PreventUserInput] priority
	 * operations. [PreventUserInput] priority should be used for operations that user input should
	 * not be able to interrupt.
	 */
	PreventUserInput,
}

/**
 * Used in place of the standard Job cancellation pathway to avoid reflective javaClass.simpleName
 * lookups to build the exception message and stack trace collection. Remove if these are changed in
 * kotlinx.coroutines.
 */
private class MutationInterruptedException : CancellationException("Mutation interrupted")

/**
 * Mutual exclusion for UI state mutation over time.
 *
 * [mutate] permits interruptible state mutation over time using a standard [MutatePriority]. A
 * [MutatorMutex] enforces that only a single writer can be active at a time for a particular state
 * resource. Instead of queueing callers that would acquire the lock like a traditional [Mutex], new
 * attempts to [mutate] the guarded state will either cancel the current mutator or if the current
 * mutator has a higher priority, the new caller will throw [CancellationException].
 *
 * [MutatorMutex] should be used for implementing hoisted state objects that many mutators may want
 * to manipulate over time such that those mutators can coordinate with one another. The
 * [MutatorMutex] instance should be hidden as an implementation detail. For example:
 */
@Stable
internal class MutatorMutex {
	private class Mutator(val priority: MutatePriority, val job: Job) {
		fun canInterrupt(other: Mutator) = priority >= other.priority

		fun cancel() = job.cancel(MutationInterruptedException())
	}

	private val currentMutator = atomicReferenceOf<Mutator?>(null)
	private val mutex = Mutex()

	private fun tryMutateOrCancel(mutator: Mutator) {
		while (true) {
			val oldMutator = currentMutator.get()
			if (oldMutator == null || mutator.canInterrupt(oldMutator)) {
				if (currentMutator.compareAndSet(oldMutator, mutator)) {
					oldMutator?.cancel()
					break
				}
			} else {
				throw CancellationException("Current mutation had a higher priority")
			}
		}
	}

	/**
	 * Enforce that only a single caller may be active at a time.
	 *
	 * If [mutate] is called while another call to [mutate] or [mutateWith] is in progress, their
	 * [priority] values are compared. If the new caller has a [priority] equal to or higher than
	 * the call in progress, the call in progress will be cancelled, throwing
	 * [CancellationException] and the new caller's [block] will be invoked. If the call in progress
	 * had a higher [priority] than the new caller, the new caller will throw
	 * [CancellationException] without invoking [block].
	 *
	 * @param priority the priority of this mutation; [MutatePriority.Default] by default. Higher
	 *   priority mutations will interrupt lower priority mutations.
	 * @param block mutation code to run mutually exclusive with any other call to [mutate] or
	 *   [mutateWith].
	 */
	suspend fun <R> mutate(
		priority: MutatePriority = MutatePriority.Default,
		block: suspend () -> R,
	) = coroutineScope {
		val mutator = Mutator(priority, coroutineContext[Job]!!)

		tryMutateOrCancel(mutator)

		mutex.withLock {
			try {
				block()
			} finally {
				currentMutator.compareAndSet(mutator, null)
			}
		}
	}

	/**
	 * Enforce that only a single caller may be active at a time.
	 *
	 * If [mutateWith] is called while another call to [mutate] or [mutateWith] is in progress,
	 * their [priority] values are compared. If the new caller has a [priority] equal to or higher
	 * than the call in progress, the call in progress will be cancelled, throwing
	 * [CancellationException] and the new caller's [block] will be invoked. If the call in progress
	 * had a higher [priority] than the new caller, the new caller will throw
	 * [CancellationException] without invoking [block].
	 *
	 * This variant of [mutate] calls its [block] with a [receiver], removing the need to create an
	 * additional capturing lambda to invoke it with a receiver object. This can be used to expose a
	 * mutable scope to the provided [block] while leaving the rest of the state object read-only.
	 * For example:
	 *
	 * @param receiver the receiver `this` that [block] will be called with
	 * @param priority the priority of this mutation; [MutatePriority.Default] by default. Higher
	 *   priority mutations will interrupt lower priority mutations.
	 * @param block mutation code to run mutually exclusive with any other call to [mutate] or
	 *   [mutateWith].
	 */
	suspend fun <T, R> mutateWith(
		receiver: T,
		priority: MutatePriority = MutatePriority.Default,
		block: suspend T.() -> R,
	) = coroutineScope {
		val mutator = Mutator(priority, coroutineContext[Job]!!)

		tryMutateOrCancel(mutator)

		mutex.withLock {
			try {
				receiver.block()
			} finally {
				currentMutator.compareAndSet(mutator, null)
			}
		}
	}
}
